use crate::ast::*;
use crate::error::{CompilerError, Result};
use super::super::super::TypeEnvironment;
use super::super::super::TypeInference;

impl TypeInference {




    pub(in crate::type_check::type_inference) fn infer_option_method_call(
        &mut self,
        option_type: &crate::ast::types::OptionType,
        method_name: &str,
        arguments: &[Expression],
        location: &crate::error::Location,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        match method_name {
            "is_some" => {
                if !arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.is_some() expects 0 arguments, got {}",
                            arguments.len()
                        ),
                    ));
                }
                Ok(Type::Basic(crate::ast::types::BasicType::Bool))
            }
            "is_none" => {
                if !arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.is_none() expects 0 arguments, got {}",
                            arguments.len()
                        ),
                    ));
                }
                Ok(Type::Basic(crate::ast::types::BasicType::Bool))
            }
            "unwrap" => {
                if !arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.unwrap() expects 0 arguments, got {}",
                            arguments.len()
                        ),
                    ));
                }
                Ok(*option_type.some_type.clone())
            }
            "expect" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.expect() expects 1 argument (message), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Check that the argument is a string
                let msg_type = self.infer_expression(&arguments[0], env)?;
                self.add_constraint(
                    &msg_type,
                    &Type::Basic(crate::ast::types::BasicType::String),
                    *location,
                )?;
                Ok(*option_type.some_type.clone())
            }
            "unwrap_or" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.unwrap_or() expects 1 argument (default value), got {}",
                            arguments.len()
                        ),
                    ));
                }
                let default_type = self.infer_expression(&arguments[0], env)?;
                self.add_constraint(&default_type, &*option_type.some_type, *location)?;
                Ok(*option_type.some_type.clone())
            }
            "unwrap_or_else" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.unwrap_or_else() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if !parameter_types.is_empty() {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.unwrap_or_else() function must take 0 arguments",
                        ));
                    }

                    // Check that the function's return type matches Some type
                    let func_return_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    self.add_constraint(&func_return_type, &*option_type.some_type, *location)?;

                    Ok(*option_type.some_type.clone())
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.unwrap_or_else() expects a function argument",
                    ))
                }
            }
            "map" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.map() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.map() function must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Some type
                    self.add_constraint(&parameter_types[0], &*option_type.some_type, *location)?;

                    // Return type is Option<U> where U is the function's return type
                    let new_some_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    Ok(Type::Option(crate::ast::types::OptionType {
                        some_type: new_some_type,
                        location: *location,
                    }))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.map() expects a function argument",
                    ))
                }
            }
            "and_then" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.and_then() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.and_then() function must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Some type
                    self.add_constraint(&parameter_types[0], &*option_type.some_type, *location)?;

                    // Return type must be an Option<U>
                    let func_return_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    match *func_return_type {
                        Type::Option(new_option_type) => Ok(Type::Option(new_option_type)),
                        // If return type is not yet resolved (e.g., from lambda body inference),
                        // we need to check if it's actually an Option by looking at the body or declared return type
                        _ => {
                            // Try to infer from the function body if it's a lambda
                            if let Expression::FunctionLiteral(func_lit) = &arguments[0] {
                                if let Some(lambda) = &func_lit.lambda {
                                    // First, check if lambda has explicit return type
                                    if let Some(declared_return) = &lambda.return_type {
                                        if let Type::Option(new_option_type) = declared_return {
                                            return Ok(Type::Option(new_option_type.clone()));
                                        }
                                    }

                                    // Otherwise, infer from lambda body
                                    let mut lambda_env = env.child();
                                    for param in &lambda.parameters {
                                        lambda_env.bind_variable(
                                            param.name.clone(),
                                            param.param_type.clone(),
                                        );
                                    }

                                    // If lambda has return type, use it for context-aware inference
                                    let expected_return = lambda.return_type.as_ref();
                                    let body_type = if let Some(_expected) = expected_return {
                                        // Infer body with expected return type for context-aware inference
                                        if let Expression::Block(block_expr) = &*lambda.body {
                                            // Convert BlockExpr to BlockStmt for infer_block_with_expected
                                            // Note: BlockExpr has statements and final_expr, BlockStmt only has statements
                                            // We need to handle final_expr separately
                                            let mut block_statements =
                                                block_expr.statements.clone();
                                            // If there's a final_expr, convert it to a return statement
                                            if let Some(final_expr) = &block_expr.final_expr {
                                                block_statements.push(
                                                    crate::ast::stmt::Statement::Return(
                                                        crate::ast::stmt::ReturnStmt {
                                                            value: Some(*final_expr.clone()),
                                                            location: block_expr.location,
                                                        },
                                                    ),
                                                );
                                            }
                                            let block = crate::ast::stmt::BlockStmt {
                                                statements: block_statements,
                                                location: block_expr.location,
                                            };
                                            // Use infer_block_statement instead of infer_block_with_expected
                                            // But we need to handle expected return type
                                            // For now, just infer the block and check the type
                                            let inferred_type = self
                                                .infer_block_statement(&block, &mut lambda_env)?;
                                            // If we have an expected return type, use it for context-aware inference
                                            if let Some(expected_type) = expected_return {
                                                self.add_constraint(
                                                    &inferred_type,
                                                    expected_type,
                                                    lambda.location,
                                                )?;
                                                // Return the expected type since we're using it for context-aware inference
                                                expected_type.clone()
                                            } else {
                                                inferred_type
                                            }
                                        } else {
                                            let inferred = self.infer_expression_with_expected(
                                                &lambda.body,
                                                &mut lambda_env,
                                                expected_return,
                                            )?;
                                            if let Some(expected_type) = expected_return {
                                                self.add_constraint(
                                                    &inferred,
                                                    expected_type,
                                                    lambda.location,
                                                )?;
                                                expected_type.clone()
                                            } else {
                                                inferred
                                            }
                                        }
                                    } else {
                                        self.infer_expression(&lambda.body, &mut lambda_env)?
                                    };

                                    // Check if body returns Option
                                    if let Type::Option(new_option_type) = body_type {
                                        Ok(Type::Option(new_option_type))
                                    } else {
                                        Err(CompilerError::type_error(
                                            location.line,
                                            location.column,
                                            "Option.and_then() function must return an Option type",
                                        ))
                                    }
                                } else {
                                    Err(CompilerError::type_error(
                                        location.line,
                                        location.column,
                                        "Option.and_then() function must return an Option type",
                                    ))
                                }
                            } else {
                                Err(CompilerError::type_error(
                                    location.line,
                                    location.column,
                                    "Option.and_then() function must return an Option type",
                                ))
                            }
                        }
                    }
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.and_then() expects a function argument",
                    ))
                }
            }
            "or_else" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.or_else() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if !parameter_types.is_empty() {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.or_else() function must take 0 arguments",
                        ));
                    }

                    // Return type must be an Option<T>
                    let func_return_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    if let Type::Option(new_option_type) = *func_return_type {
                        // Check that the return type matches
                        self.add_constraint(
                            &new_option_type.some_type,
                            &*option_type.some_type,
                            *location,
                        )?;
                        Ok(Type::Option(new_option_type))
                    } else {
                        Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.or_else() function must return an Option type",
                        ))
                    }
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.or_else() expects a function argument",
                    ))
                }
            }
            "filter" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.filter() expects 1 argument (predicate function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.filter() predicate must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Some type
                    self.add_constraint(&parameter_types[0], &*option_type.some_type, *location)?;

                    // Check that the function returns bool
                    let func_return_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    self.add_constraint(
                        &func_return_type,
                        &Type::Basic(crate::ast::types::BasicType::Bool),
                        *location,
                    )?;

                    // Return type is Option<T> (same as input)
                    Ok(Type::Option(option_type.clone()))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.filter() expects a function argument",
                    ))
                }
            }
            "and" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.and() expects 1 argument (other Option), got {}",
                            arguments.len()
                        ),
                    ));
                }
                let other_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Option(other_option_type) = other_type {
                    Ok(Type::Option(other_option_type))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.and() expects an Option argument",
                    ))
                }
            }
            "or" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.or() expects 1 argument (other Option), got {}",
                            arguments.len()
                        ),
                    ));
                }
                let other_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Option(other_option_type) = other_type {
                    // Check that the other Option's Some type matches
                    self.add_constraint(
                        &other_option_type.some_type,
                        &*option_type.some_type,
                        *location,
                    )?;
                    Ok(Type::Option(crate::ast::types::OptionType {
                        some_type: option_type.some_type.clone(),
                        location: *location,
                    }))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.or() expects an Option argument",
                    ))
                }
            }
            "xor" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.xor() expects 1 argument (other Option), got {}",
                            arguments.len()
                        ),
                    ));
                }
                let other_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Option(other_option_type) = other_type {
                    // Check that the other Option's Some type matches
                    self.add_constraint(
                        &other_option_type.some_type,
                        &*option_type.some_type,
                        *location,
                    )?;
                    Ok(Type::Option(crate::ast::types::OptionType {
                        some_type: option_type.some_type.clone(),
                        location: *location,
                    }))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.xor() expects an Option argument",
                    ))
                }
            }
            "ok_or" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.ok_or() expects 1 argument (error value), got {}",
                            arguments.len()
                        ),
                    ));
                }
                let err_type = self.infer_expression(&arguments[0], env)?;
                Ok(Type::Result(crate::ast::types::ResultType {
                    ok_type: option_type.some_type.clone(),
                    err_type: Box::new(err_type),
                    location: *location,
                }))
            }
            "ok_or_else" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.ok_or_else() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if !parameter_types.is_empty() {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.ok_or_else() function must take 0 arguments",
                        ));
                    }

                    // Function's return type becomes the Err type
                    let err_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    Ok(Type::Result(crate::ast::types::ResultType {
                        ok_type: option_type.some_type.clone(),
                        err_type,
                        location: *location,
                    }))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.ok_or_else() expects a function argument",
                    ))
                }
            }
            "map_or" => {
                if arguments.len() != 2 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Option.map_or() expects 2 arguments (default, function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer default type
                let default_type = self.infer_expression(&arguments[0], env)?;

                // Infer the function type
                let func_type = self.infer_expression(&arguments[1], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.map_or() function must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Some type
                    self.add_constraint(&parameter_types[0], &*option_type.some_type, *location)?;

                    // Check that the function's return type matches default type
                    let func_return_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    self.add_constraint(&func_return_type, &default_type, *location)?;

                    // Return type is the default type (or function return type, they should match)
                    Ok(default_type)
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.map_or() expects a function as second argument",
                    ))
                }
            }
            "map_or_else" => {
                if arguments.len() != 2 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!("Option.map_or_else() expects 2 arguments (default_fn, function), got {}", arguments.len()),
                    ));
                }
                // Infer default_fn type
                let default_fn_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types: default_params,
                    return_type: default_return,
                    ..
                }) = default_fn_type
                {
                    if !default_params.is_empty() {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.map_or_else() default_fn must take 0 arguments",
                        ));
                    }
                    let default_type = default_return
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));

                    // Infer the function type
                    let func_type = self.infer_expression(&arguments[1], env)?;
                    if let Type::Function(FunctionType {
                        parameter_types,
                        return_type,
                        ..
                    }) = func_type
                    {
                        if parameter_types.len() != 1 {
                            return Err(CompilerError::type_error(
                                location.line,
                                location.column,
                                "Option.map_or_else() function must take exactly 1 argument",
                            ));
                        }
                        // Check that the function's parameter type matches Some type
                        self.add_constraint(
                            &parameter_types[0],
                            &*option_type.some_type,
                            *location,
                        )?;

                        // Check that the function's return type matches default type
                        let func_return_type = return_type
                            .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                        self.add_constraint(&func_return_type, &*default_type, *location)?;

                        // Return type is the default type (or function return type, they should match)
                        Ok(*default_type)
                    } else {
                        Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Option.map_or_else() expects a function as second argument",
                        ))
                    }
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Option.map_or_else() expects a function as first argument",
                    ))
                }
            }
            _ => Err(CompilerError::type_error(
                location.line,
                location.column,
                format!("Unknown Option method: {}", method_name),
            )),
        }
    }
}
